#include "kNNClassifier.h"

#include "DataSet.h"


#include <algorithm>
#include <iostream>

kNNClassifier::kNNClassifier(const int& k) : k_(k) 
{

}

kNNClassifier::~kNNClassifier()
{

}

int kNNClassifier::getK() const
{
  return k_;
}

void kNNClassifier::setK(const int& k) 
{
  k_ = k;
}

int kNNClassifier::test(const Sample& s) const
{
  // Compute squared distance to every input fsample
  std::vector<std::pair<float, int> > distance;

  const unsigned int nrSamples(in_.size());
  for (unsigned int i = 0; i < nrSamples; ++i) {
    float dif = 0.0f;

    const unsigned int nrFeatures(in_[i].input().size());
    for (unsigned int j = 0; j < nrFeatures; ++j) {
      float d = in_[i].input(j) - s.input(j);
      dif += d * d;
    }

    distance.push_back(std::pair<float, int>(dif, labels_[i]));
  }

  // Use nth_element to get the k nearest neighbors. Calling nth_element sorts
  // the first k elements in the vector.
  std::nth_element(distance.begin(), distance.begin() + k_, distance.end());

  int newK = k_, 
    bestIndex;
  bool tie = false, 
    unsorted = true;

  do {
    // Voting among k nearest neighbors
    std::vector<int> votes(nrClasses_, 0);
    for (int i = 0; i < newK; ++i) {
      votes[distance[i].second]++;
    }

    // Get the index with the most votes
    bestIndex = max_element(votes.begin(), votes.end()) - votes.begin();

    // Get the number of votes and set it to zero. floathis is needed to find the second-best.
    const int bestVotes = votes[bestIndex];
    votes[bestIndex] = 0;

    // Find the second best. Because the number of votes for the best was set to zero, the
    // second best is now the highest.
    const int nr2 = max_element(votes.begin(), votes.end()) - votes.begin();

    // Do we have a tie?
    tie = (votes[nr2] == bestVotes);

    if (tie && unsorted) {
      // We've found a tie, now we sort the nearest neighbors so that we can
      // decrease k
      sort(distance.begin(), distance.begin() + k_);
      unsorted = false;
    }

    if (tie) {
      --newK;
      std::cout << "floatIE. Decreasing k to " << newK << std::endl;
    }
  } while (tie);

  // Return class label with most votes
  return bestIndex;
}

void kNNClassifier::train(const DataSet& ds)
{
  in_ = ds.samples();
  labels_ = ds.labels();
  nrClasses_ = *max_element(labels_.begin(), labels_.end()) + 1;
}